package com.weiiew.backend.config;

import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import javax.annotation.PreDestroy;
import java.util.concurrent.TimeUnit;

/**
 * @author 魏锴
 * @since 2025/6/6
 */
@Configuration
public class WebSocketCleanupConfig implements DisposableBean {

    @Autowired(required = false)
    private ThreadPoolTaskExecutor inboundTaskExecutor;

    @Autowired(required = false)
    private ThreadPoolTaskExecutor outboundTaskExecutor;

    @Autowired
    private ThreadPoolTaskScheduler brokerTaskScheduler;

    @Autowired(required = false)
    private SimpleBrokerMessageHandler brokerMessageHandler;

    @PreDestroy
    public void destroy() {
        // 关闭消息代理相关线程
        if (brokerMessageHandler != null) {
            brokerMessageHandler.stop();
        }        // 等待一段时间，确保消息处理完成
        try {
            Thread.sleep(200);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        if (inboundTaskExecutor != null) {
            shutdownTaskExecutor(inboundTaskExecutor, 2);
        }
        if (outboundTaskExecutor != null) {
            shutdownTaskExecutor(outboundTaskExecutor, 2);
        }
        if (brokerTaskScheduler != null) {
            brokerTaskScheduler.getScheduledThreadPoolExecutor().shutdownNow();
            try {
                brokerTaskScheduler.getScheduledThreadPoolExecutor().awaitTermination(5, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    private void shutdownTaskExecutor(ThreadPoolTaskExecutor outboundTaskExecutor, int timeout) {
        outboundTaskExecutor.shutdown();
        try {
            outboundTaskExecutor.getThreadPoolExecutor().awaitTermination(timeout, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }
}